#!/usr/bin/env python3
"""
简易Python转SO工具
功能：将指定Python文件编译为.so模块，输出到当前目录的compiled文件夹
"""

import os
import sys
import shutil
import subprocess
import tempfile
import re
from pathlib import Path

# ==================== 配置区域 (在此处修改设置) ====================
# 要编译的Python文件列表（可以是相对路径或绝对路径）
PY_FILES = [
    "a1.py",   # 示例：不合法的文件名（以数字开头）
    "a3.py",   # 示例：不合法的文件名
    # "module1.py",  # 合法的文件名（字母开头）
    # 可以添加更多文件...
]

# 是否保留源文件 (True=保留, False=删除)
KEEP_SOURCE = True

# 输出目录名称（将创建在当前目录下）
OUTPUT_DIR = "compiled"
# ==================================================================

def is_valid_module_name(name):
    """检查文件名是否符合Python模块命名规范"""
    # 模块名规则：只能包含字母、数字、下划线，且不能以数字开头
    return re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', name) is not None

def compile_py_to_so(py_file, output_dir, keep_source):
    """编译单个Python文件为.so模块"""
    try:
        py_path = Path(py_file).resolve()
        if not py_path.exists():
            print(f"错误: 文件不存在 - {py_path}")
            return False

        if py_path.suffix != ".py":
            print(f"错误: 不是Python文件 - {py_path}")
            return False

        # 获取文件名（不含扩展名）
        module_name = py_path.stem
        
        # 检查模块名合法性
        if not is_valid_module_name(module_name):
            print(f"错误: 文件名'{module_name}'不符合Python模块命名规范")
            print("      提示：模块名只能包含字母、数字、下划线，且不能以数字开头")
            return False

        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        # 创建临时目录用于编译
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_dir = Path(tmp_dir)
            tmp_py = tmp_dir / f"{module_name}.py"
            tmp_setup = tmp_dir / "setup.py"

            # 复制源文件到临时目录
            shutil.copy(py_path, tmp_py)

            # 生成setup.py
            setup_code = f'''
from setuptools import setup
from Cython.Build import cythonize

setup(
    ext_modules=cythonize(
        "{module_name}.py",
        compiler_directives={{"language_level": 3}}
    ),
    script_args=["build_ext", "--inplace"]
)
'''
            tmp_setup.write_text(setup_code)

            # 执行编译
            print(f"正在编译: {py_path.name}")
            result = subprocess.run(
                [sys.executable, str(tmp_setup)],
                cwd=tmp_dir,
                capture_output=True,
                text=True
            )

            if result.returncode != 0:
                print(f"编译失败: {py_path.name}")
                print(result.stderr)
                return False

            # 查找并移动生成的.so文件
            so_files = list(tmp_dir.glob(f"{module_name}.cpython-*.so"))
            if not so_files:
                print(f"未找到生成的.so文件: {py_path.name}")
                return False

            so_file = so_files[0]
            target_so = output_dir / so_file.name
            shutil.move(str(so_file), target_so)
            print(f"已生成: {target_so}")

        # 处理源文件
        if not keep_source:
            py_path.unlink()
            print(f"已删除源文件: {py_path.name}")

        return True

    except Exception as e:
        print(f"处理{py_file}时出错: {str(e)}")
        return False

def main():
    print("===== 简易Python转SO工具 =====")
    print(f"输出目录: {os.path.abspath(OUTPUT_DIR)}")
    print(f"保留源文件: {'是' if KEEP_SOURCE else '否'}")
    print(f"待编译文件: {len(PY_FILES)}个")
    print("----------------------------")

    success = 0
    for file in PY_FILES:
        if compile_py_to_so(file, OUTPUT_DIR, KEEP_SOURCE):
            success += 1

    print("\n===== 编译完成 =====")
    print(f"成功: {success}/{len(PY_FILES)}")
    print(f"输出文件位于: {os.path.abspath(OUTPUT_DIR)}")

if __name__ == "__main__":
    main()
